# -*- coding: UTF-8 -*-
import random
import numpy as np
import pickle
import sys
import shutil, os

# 将上级目录加入系统路径，以便导入上级目录中的模块
sys.path.append('..')
# from data_loc.draw_gps import draw_gps  # 该行被注释，可能用于绘制GPS轨迹
import config
from taxi_special import square  # 从 taxi_special 模块导入 square 函数

WAIT = 20  # 等待时间，用于扩展短路线的数据

class GenerateUser:
    """
    生成用户轨迹数据的类
    """
    def __init__(self):
        """
        初始化 taxi_mob_all 用于存储所有出租车的移动数据
        """
        self.taxi_mob_all = {}

    def ini_loc(self, curseed):
        """
        初始化用户位置，从出租车数据集中随机选择用户轨迹
        :param curseed: 随机种子，确保每次生成的数据一致
        :return: 选中的出租车轨迹数据
        """
        random.seed(curseed)  # 设置随机种子，保证结果可复现
        absolute_path2 = "taxi_npy/"  # 存储出租车数据的路径
        
        # 读取所有的出租车 ID
        with open(absolute_path2 + 'taxi_id.pkl', 'rb') as fu:
            self.new_taxi_ids = pickle.load(fu)
        
        # 随机选择符合 config.USER_NUM 数量的出租车 ID 作为用户
        tmp_list = list(self.new_taxi_ids)
        select_ids = random.sample(list(tmp_list), config.USER_NUM) 
        
        # 读取所选出租车的轨迹数据
        for idx, item2 in enumerate(select_ids):
            detail_path = str(item2) + '.pkl'
            with open(absolute_path2 + detail_path, 'rb') as fu2:
                one_taxi_mob_list = pickle.load(fu2)
            self.taxi_mob_all.update({idx: one_taxi_mob_list})
        
        return self.taxi_mob_all

class User:
    """
    用户类，包含用户的基本信息和任务
    """
    def __init__(self, uid, ulocx, ulocy, utskids, ulayernamelist, tsk_arrive_time, replace) -> None:
        """
        初始化用户对象
        :param uid: 用户 ID
        :param ulocx: 用户当前位置的 x 坐标
        :param ulocy: 用户当前位置的 y 坐标
        :param utskids: 用户拥有的任务 ID 列表
        :param ulayernamelist: 用户所属的层列表
        :param tsk_arrive_time: 任务到达时间
        :param replace: 用户 ID 替换值
        """
        self.uid = uid
        self.user_locationx = ulocx
        self.user_locationy = ulocy
        self.usr_has_tsk = utskids  # 用户拥有的任务 ID
        self.usr_has_lay = ulayernamelist  # 用户所属层列表
        self.tsk_arrive_time_start = tsk_arrive_time  # 任务到达时间
        self.usrid_replace = replace  # 用户 ID 替换值


def generate_total_trace():
    """
    生成所有用户的轨迹数据，提取符合区域范围的出租车轨迹
    """
    lat_min, lon_min, lat_max, lon_max = square()  # 获取轨迹范围
    absolute_path = '../taxi/'  # 出租车数据存储路径
    root_path = '../taxi/'
    taxi_file = '_taxi_id.txt'  # 存储出租车 ID 的文件
    
    all_taxi_id = []  # 存储所有出租车 ID
    with open(absolute_path + taxi_file, 'r') as f:
        for line in f:
            all_taxi_id.append(line.strip())
    
    taxi_num = len(all_taxi_id)
    print("task number is:" + str(taxi_num))

    new_taxi_id = []  # 存储符合范围的出租车 ID
    new_taxi_mob = {}  # 存储符合范围的出租车轨迹
    
    for id in all_taxi_id:
        detail_path = str(id) + '.txt'
        with open(absolute_path + detail_path, 'r') as f2:
            one_taxi = f2.readlines()
        
        one_taxi_mob = []  # 存储单辆出租车的轨迹数据
        tmp_step = 0
        
        for idx, line in enumerate(one_taxi[:-1]):
            lat = eval(line.strip().split()[0])
            lon = eval(line.strip().split()[1])
            
            # 判断是否超出范围
            if lon <= lon_min or lon >= lon_max or lat <= lat_min or lat >= lat_max:
                break
            
            one_taxi_mob.append([lat, lon])
            tmp_step += 1
            
            seconds = eval(line.strip().split()[-1])
            seconds_next = eval(one_taxi[idx+1].strip().split()[-1])
            dd = seconds_next - seconds
            while dd > 0:
                one_taxi_mob.append([lat, lon])
                tmp_step += 1
                dd -= 1
        
        if len(one_taxi_mob) > 1000:  # 轨迹点数大于1000的才保留
            if len(one_taxi_mob) < 10000:
                new_taxi_id.append(eval(id))
                tmp_one_taxi_mob = []
                for item in one_taxi_mob:
                    p = WAIT
                    while p > 0:
                        tmp_one_taxi_mob.append(item)
                        p -= 1
                new_taxi_mob.update({eval(id): tmp_one_taxi_mob})
            else:
                new_taxi_id.append(eval(id))
                new_taxi_mob.update({eval(id): one_taxi_mob})
    
    # 保存符合范围的轨迹数据
    absolute_path2 = 'taxi_npy/'
    id_new = []
    
    for idx, taxi_id in enumerate(new_taxi_id):
        tmp_mob = new_taxi_mob.get(taxi_id, None)
        path_mob = str(idx) + '.pkl'
        with open(absolute_path2 + path_mob, 'wb') as ff2:
            pickle.dump(tmp_mob, ff2)
        id_new.append(idx)
    
    with open(absolute_path2 + 'taxi_id.pkl', 'wb') as ff1:
        pickle.dump(id_new, ff1)
    print('get-NUM: ', len(id_new))
                
if __name__ == '__main__':
    generate_total_trace()  # 运行轨迹生成函数